import os
import argparse
import xml.etree.ElementTree as ET
from PIL import Image
import shutil

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--classes_file', type=str, required=True)
    parser.add_argument('--txt_folder1', type=str, required=True)
    parser.add_argument('--txt_folder2', type=str, required=True)
    parser.add_argument('--img_folder', type=str, required=True)
    parser.add_argument('--same_folder', type=str, required=True)
    parser.add_argument('--not_same_folder', type=str, required=True)
    parser.add_argument('--iou_threshold', type=float, required=True)
    return parser.parse_args()

def load_classes(classes_file):
    with open(classes_file, 'r') as f:
        classes = [line.strip() for line in f.readlines()]
    return classes

def load_txt_annotations(txt_folder):
    annotations = {}
    for file in os.listdir(txt_folder):
        if file.endswith('.txt'):
            img_name = file[:-4]
            annotations[img_name] = []
            with open(os.path.join(txt_folder, file), 'r') as f:
                for line in f.readlines():
                    class_id, x, y, w, h = line.strip().split()
                    annotations[img_name].append({
                        'class_id': int(class_id),
                        'x': float(x),
                        'y': float(y),
                        'w': float(w),
                        'h': float(h)
                    })
    return annotations

def calculate_iou(box1, box2):
    x1, y1, w1, h1 = box1['x'], box1['y'], box1['w'], box1['h']
    x2, y2, w2, h2 = box2['x'], box2['y'], box2['w'], box2['h']
    xi1, yi1, xi2, yi2 = max(x1, x2), max(y1, y2), min(x1 + w1, x2 + w2), min(y1 + h1, y2 + h2)
    inter_area = max(0, xi2 - xi1) * max(0, yi2 - yi1)
    box1_area = w1 * h1
    box2_area = w2 * h2
    iou = inter_area / (box1_area + box2_area - inter_area)
    return iou

def create_xml_annotation(img_name, annotations, classes, img_folder, xml_folder):
    img_path = os.path.join(img_folder, img_name + '.jpg')
    img_file = Image.open(img_path)
    width, height = img_file.size
    xml_file = open(os.path.join(xml_folder, img_name + '.xml'), 'w')
    xml_file.write('<?xml version="1.0"?>\n')
    xml_file.write('<annotation>\n')
    xml_file.write('\t<folder>simple</folder>\n')
    xml_file.write('\t<filename>' + img_name + '.jpg</filename>\n')
    xml_file.write('\t<size>\n')
    xml_file.write('\t\t<width>' + str(width) + '</width>\n')
    xml_file.write('\t\t<height>' + str(height) + '</height>\n')
    xml_file.write('\t\t<depth>3</depth>\n')
    xml_file.write('\t</size>\n')
    for annotation in annotations:
        class_id = annotation['class_id']
        class_name = classes[class_id]
        x_center = annotation['x']
        y_center = annotation['y']
        w = annotation['w']
        h = annotation['h']
        xmax = int((2 * x_center * width + w * width) / 2)
        xmin = int((2 * x_center * width - w * width) / 2)
        ymax = int((2 * y_center * height + h * height) / 2)
        ymin = int((2 * y_center * height - h * height) / 2)
        xml_file.write('\t<object>\n')
        xml_file.write('\t\t<name>' + class_name + '</name>\n')
        xml_file.write('\t\t<pose>Unspecified</pose>\n')
        xml_file.write('\t\t<truncated>0</truncated>\n')
        xml_file.write('\t\t<difficult>0</difficult>\n')
        xml_file.write('\t\t<bndbox>\n')
        xml_file.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
        xml_file.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
        xml_file.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
        xml_file.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
        xml_file.write('\t\t</bndbox>\n')
        xml_file.write('\t</object>\n')
    xml_file.write('</annotation>')

def main():
    args = parse_args()
    classes = load_classes(args.classes_file)
    txt_annotations1 = load_txt_annotations(args.txt_folder1)
    txt_annotations2 = load_txt_annotations(args.txt_folder2)

    same_folder = args.same_folder
    not_same_folder = args.not_same_folder
    iou_threshold = args.iou_threshold

    os.makedirs(same_folder, exist_ok=True)
    os.makedirs(not_same_folder, exist_ok=True)
    os.makedirs(os.path.join(same_folder, 'img_same'), exist_ok=True)
    os.makedirs(os.path.join(same_folder, 'xml_same'), exist_ok=True)
    os.makedirs(os.path.join(not_same_folder, 'img_not_same'), exist_ok=True)
    os.makedirs(os.path.join(not_same_folder, 'xml_not_same'), exist_ok=True)

    img_names = [os.path.splitext(file)[0] for file in os.listdir(args.img_folder) if file.endswith('.jpg')]
    for img_name in img_names:
        annotations1 = txt_annotations1.get(img_name, [])
        annotations2 = txt_annotations2.get(img_name, [])

        if not annotations1 and not annotations2:
            # Create an empty XML file
            create_xml_annotation(img_name, [], classes, args.img_folder, os.path.join(not_same_folder, 'xml_not_same'))
            shutil.move(os.path.join(args.img_folder, img_name + '.jpg'), os.path.join(os.path.join(not_same_folder, 'img_not_same'), img_name + '.jpg'))
        elif len(annotations1) != len(annotations2):
            if annotations1:
                create_xml_annotation(img_name, annotations1, classes, args.img_folder, os.path.join(not_same_folder, 'xml_not_same'))
            else:
                create_xml_annotation(img_name, annotations2, classes, args.img_folder, os.path.join(not_same_folder, 'xml_not_same'))
            shutil.move(os.path.join(args.img_folder, img_name + '.jpg'), os.path.join(os.path.join(not_same_folder, 'img_not_same'), img_name + '.jpg'))
        else:
            same = True
            for annotation1, annotation2 in zip(annotations1, annotations2):
                if annotation1['class_id'] != annotation2['class_id']:
                    same = False
                    break
                iou = calculate_iou(annotation1, annotation2)
                if iou < iou_threshold:
                    same = False
                    break
            if same:
                create_xml_annotation(img_name, annotations1, classes, args.img_folder, os.path.join(same_folder, 'xml_same'))
                shutil.move(os.path.join(args.img_folder, img_name + '.jpg'), os.path.join(os.path.join(same_folder, 'img_same'), img_name + '.jpg'))
            else:
                create_xml_annotation(img_name, annotations1, classes, args.img_folder, os.path.join(not_same_folder, 'xml_not_same'))
                shutil.move(os.path.join(args.img_folder, img_name + '.jpg'), os.path.join(os.path.join(not_same_folder, 'img_not_same'), img_name + '.jpg'))

if __name__ == '__main__':
    main()

# python txt_to_xml_2model.py --classes_file class.names --txt_folder1 /home/chenwh/yolov8_rknn_ultralytics-8.2.2/runs/detect/predict/labels --txt_folder2 /home/chenwh/yolov8_rknn_ultralytics-8.2.2/runs/detect/predict2/labels --img_folder traindata_gt_guiji_old_960_540_5 --same_folder traindata_gt_guiji_old_960_540_5_same --not_same_folder traindata_gt_guiji_old_960_540_5_not_same --iou_threshold 0.8
